{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from os.path import join\n",
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "import pdb, sys, torch\n",
    "sys.path.append('.')\n",
    "torch.backends.cuda.matmul.allow_tf32 = True\n",
    "torch.backends.cudnn.allow_tf32 = True\n",
    "torch.backends.cudnn.deterministic = True\n",
    "from diffuser.guides.policies import Policy\n",
    "import diffuser.datasets as datasets\n",
    "import diffuser.utils as utils\n",
    "from diffuser.models import GaussianDiffusionPB\n",
    "from datetime import datetime\n",
    "import os.path as osp\n",
    "from importlib import reload\n",
    "from diffuser.utils.jupyter_utils import suppress_stdout\n",
    "np.set_printoptions(precision=3)\n",
    "\n",
    "## You need to download 'kuka7d-base' from OneDrive Link in README.md to launch\n",
    "class Parser(utils.Parser):\n",
    "    config: str = \"config/kuka7d/kuka_exp.py\"\n",
    "\n",
    "#---------------------------------- setup ----------------------------------#\n",
    "\n",
    "## training args\n",
    "args_train = Parser().parse_args('diffusion', from_jupyter=True)\n",
    "args = Parser().parse_args('plan', from_jupyter=True)\n",
    "\n",
    "args.savepath = None # osp.join(args.savepath, sub_dir)\n",
    "args.load_unseen_maze = True\n",
    "\n",
    "## load dataset here, dataset is a string: name of the env\n",
    "print('args.dataset', type(args.dataset), args.dataset)\n",
    "print('args.dataset_eval', type(args.dataset_eval), args.dataset_eval)\n",
    "use_normed_wallLoc = args_train.dataset_config.get('use_normed_wallLoc', False)\n",
    "\n",
    "## use the trained env or eval env\n",
    "load_unseen_maze = args.load_unseen_maze # True False\n",
    "\n",
    "with suppress_stdout():\n",
    "    # ---------- load normalizer ------------\n",
    "    train_env_list = datasets.load_environment(args.dataset, is_eval=load_unseen_maze)\n",
    "    train_normalizer = utils.load_datasetNormalizer(train_env_list.dataset_url, args_train, train_env_list)\n",
    "\n",
    "    #---------------------------------- loading ----------------------------------#\n",
    "    ld_config = dict(env_instance=train_env_list) \n",
    "    ## dataset, renderer, model, diffusion, trainer.ema_model, trainer, epoch\n",
    "    diffusion_experiment = utils.load_potential_diffusion_model(args.logbase, args.dataset, \\\n",
    "                args_train.exp_name, epoch=args.diffusion_epoch, ld_config=ld_config)\n",
    "    diffusion = diffusion_experiment.ema"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### load eval problems"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffuser.guides.kuka_plan_utils import load_eval_problems_pb\n",
    "problems_dict = load_eval_problems_pb(train_env_list, {'load_unseen_maze': load_unseen_maze})\n",
    "probs: np.ndarray\n",
    "probs = problems_dict['problems'] # (300, 200, 2, 7)\n",
    "print('loaded epoch:', diffusion_experiment.epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import diffuser.utils.jupyter_utils as vu; reload(vu)\n",
    "from diffuser.utils.jupyter_utils import get_all_suc_trajs\n",
    "print('condition_guidance_w', diffusion.condition_guidance_w )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## example of creating an env and get the wall locations\n",
    "env_single = train_env_list.create_single_env(17)\n",
    "env_single.wall_locations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diffusion: GaussianDiffusionPB\n",
    "bs = 100\n",
    "policy = Policy(diffusion, train_normalizer, use_ddim=True)\n",
    "test_hor = 48\n",
    "diffusion.horizon = test_hor\n",
    "\n",
    "wloc_source = 'dataset'\n",
    "if wloc_source == 'dataset':\n",
    "    ## env \n",
    "    env_id, p_id = 87, 8\n",
    "\n",
    "    env_single = train_env_list.create_single_env(env_id)\n",
    "    prob = probs[env_id, p_id, ] # (2, 7)\n",
    "    ## env \n",
    "    start = prob[None, 0].repeat( bs, 0 ) # (bs, 7)\n",
    "    goal =  prob[None, 1].repeat( bs, 0 )\n",
    "    \n",
    "elif wloc_source == 'custom':\n",
    "    '''Define your own wall locations'''\n",
    "    env_id = 0 ## will overwrite env0\n",
    "    xyz_list = np.array([\n",
    "        [-0.492,  0.177,  0.846],\n",
    "        [-0.045, -0.433,  0.291],\n",
    "        [ 0.431,  0.476,  0.409],\n",
    "        [ 0.469, -0.03 ,  0.886]])\n",
    "    assert len(xyz_list) == 4, 'since w/o using composition, should equal to training number.'\n",
    "    \n",
    "    env_single = train_env_list.create_env_by_pos(env_id, xyz_list)\n",
    "    ## define the start and goal joint state as well...\n",
    "\n",
    "## prepare wloc\n",
    "wloc_np = train_env_list.wallLoc_list[env_id]\n",
    "wloc_tensor = utils.to_torch(wloc_np).reshape(1, -1).repeat( (bs, 1) )\n",
    "print( f'wloc_np: ', wloc_np.shape )\n",
    "\n",
    "cond = {0: start, \n",
    "        test_hor-1: goal}\n",
    "\n",
    "samples = policy(cond, batch_size=-1, wall_locations=wloc_tensor, use_normed_wallLoc=use_normed_wallLoc, return_diffusion=False)\n",
    "## samples[0] is dummy, can be ignored\n",
    "unnm_traj = samples[1].observations # 3d (B, H, 7)\n",
    "\n",
    "unnm_suc_traj, suc_idxs = get_all_suc_trajs( env_single, unnm_traj, goal=goal)\n",
    "num_suc_trajs = len(unnm_suc_traj)\n",
    "print(f'num_suc_trajs: {num_suc_trajs}') # a batch of traj\n",
    "print(f'unnm_traj: {len(unnm_traj)}, {unnm_traj[0].shape}') # a batch of traj\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('start:', np.array2string(start[0], separator=', '))\n",
    "print('goal:',  np.array2string(goal[0], separator=', '))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Do visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h = w = 300 # render resolution\n",
    "pfr = 77\n",
    "import pybullet as p\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm; \n",
    "from pb_diff_envs.utils.kuka_utils_luo import add_start_end_marker\n",
    "\n",
    "'''\n",
    "generate kuka video for website and ppt slides\n",
    "'''\n",
    "def visualize_kuka_traj_luo_web(env, traj, lock=None, is_ee3d=False, is_debug=False, is_mpo=False):\n",
    "    '''\n",
    "    The difference with visualize_traj is that this function directly receives \n",
    "    [a list of 1d numpy] or [2d np] as trajectory, not using a traj_agent\n",
    "    \n",
    "    '''\n",
    "    ## the objects' trajectories, should be 2 for static obj\n",
    "    max_len_traj = max([len(obj.trajectory.waypoints) for obj in env.objects])\n",
    "    gifs = []\n",
    "    max_len = max(len(traj), max_len_traj)\n",
    "    assert max_len == len(traj)\n",
    "    # print('[vis traj] max_len_traj', max_len_traj)\n",
    "    print('[vis traj] max_len:', max_len)\n",
    "\n",
    "    # se_valid = check_start_end(traj) # if the start and end is valid\n",
    "    se_valid = None\n",
    "    # if not p.isConnected():\n",
    "    if lock is not None:\n",
    "        lock.acquire()\n",
    "        env.load(GUI=False)\n",
    "        lock.release()\n",
    "    else:\n",
    "        # if p.isConnected()\n",
    "        # print(f'p isConnected: {p.isConnected()}')\n",
    "        env.load(GUI=is_debug) # ok\n",
    "\n",
    "    i_y, i_p = 0, 0\n",
    "    colors = plt.cm.jet(np.linspace(0,1, max_len)) # color of the dot\n",
    "    vshape_id = p.createVisualShape(p.GEOM_SPHERE, radius=0.02)  # Red color\n",
    "    has_collision = False\n",
    "    num_collision_frame = 0\n",
    "    collision_list = []\n",
    "    add_start_end_marker(env, traj)\n",
    "    if is_mpo:\n",
    "        traj = traj[:-1]; max_len -= 1\n",
    "\n",
    "    assert not is_ee3d, 'depricated, prevent bugs'\n",
    "    # for c, timestep in tqdm(enumerate(traj)):\n",
    "    for i_t in tqdm(range(max_len)):\n",
    "\n",
    "        if not is_ee3d: # traj is xyz-level\n",
    "            env.robot.set_config(traj[i_t])\n",
    "            new_pos = p.getLinkState(env.robot_id, 6)[0]\n",
    "        else:\n",
    "            new_pos = traj[i_t]\n",
    "        cam_pos = p.getLinkState(env.robot_id, 5)[0]\n",
    "        new_pos_q = p.getQuaternionFromEuler(new_pos)\n",
    "\n",
    "        dist = 2 - np.linalg.norm(cam_pos)\n",
    "        # tmp_id = p.loadURDF(\"sphere2red.urdf\", new_pos, globalScaling=0.05, flags=p.URDF_IGNORE_COLLISION_SHAPES)\n",
    "        ## slow\n",
    "        tmp_id = p.createMultiBody(baseMass=0, basePosition=new_pos, baseVisualShapeIndex=vshape_id,)\n",
    "        p.changeVisualShape(tmp_id, linkIndex=-1, rgbaColor=colors[i_t])\n",
    "\n",
    "        # print(p.getBodyInfo(tmp_id))\n",
    "        p.performCollisionDetection()\n",
    "        c_pts = p.getContactPoints(env.robot_id)\n",
    "        if c_pts is None:\n",
    "            pass # also no collision\n",
    "        elif len(c_pts) > 0: # very important, check is None\n",
    "            has_collision = True\n",
    "            num_collision_frame += 1\n",
    "            collision_list.append(i_t)\n",
    "            # print(f'has_collision {has_collision}; c_pts: {c_pts}')\n",
    "\n",
    "        ## -----------------------------------\n",
    "        ## Configurate Camera\n",
    "        cam_pos = [0,0,0.3]\n",
    "        dist = 1.95 # original exp settings:1.8\n",
    "        \n",
    "        view_mat = p.computeViewMatrixFromYawPitchRoll(\n",
    "                cameraTargetPosition=cam_pos, distance=dist, # [0, 0, 0]\n",
    "                yaw=-i_y, pitch=-30-i_p, roll=0, upAxisIndex=2) # yaw=-90-i_y\n",
    "        i_p = i_p + 0.2 if i_p < 20 else i_p\n",
    "        i_y += 0.3\n",
    "\n",
    "        ## ---- for web high-quality visualization ------\n",
    "\n",
    "        resol_mul = 1\n",
    "        k7d_camera_info = (h, w, (0.886519193649292, 0.2863903045654297, -0.36340656876564026, 0.0, -0.4626918435096741, 0.5487248301506042, -0.6962882280349731, 0.0, 0.0, 0.7854180932044983, 0.6189655065536499, 0.0, -0.0, 1.1920928955078125e-07, -2.101985216140747, 1.0), (1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, -1.0000200271606445, -1.0, 0.0, 0.0, -0.02000020071864128, 0.0), (0.0, 0.0, 1.0), (0.36340656876564026, 0.6962882280349731, -0.6189655065536499), (17730.3828125, -9253.8369140625, 0.0), (5727.80615234375, 10974.49609375, 15708.36328125), -27.56093978881836, -38.24062728881836, 2.101985216140747, (-0.0, -0.0, 0.0))\n",
    "                           \n",
    "\n",
    "        act_yaw, act_pitch, act_dist, act_target = k7d_camera_info[-4:]\n",
    "        \n",
    "        view_mat = k7d_camera_info[2]\n",
    "\n",
    "        width, height, img, depth, seg = p.getCameraImage(\n",
    "            width=k7d_camera_info[0] * resol_mul,  # Image width\n",
    "            height=k7d_camera_info[1] * resol_mul,  # Image height\n",
    "            # viewMatrix=k7d_camera_info[2],  # View matrix\n",
    "            viewMatrix=view_mat,  # View matrix\n",
    "            projectionMatrix=k7d_camera_info[3],  # Projection matrix\n",
    "            flags=p.ER_NO_SEGMENTATION_MASK,\n",
    "            renderer=p.ER_BULLET_HARDWARE_OPENGL,\n",
    "        )\n",
    "        gifs.append( img )\n",
    "\n",
    "        ## ----------------------------------\n",
    "    \n",
    "    ## stop at the last one for 2 sec\n",
    "    ds = [150] * (len(gifs)-1)\n",
    "    ds.append(2000)\n",
    "\n",
    "    def process_elem(st, end, n_interp, idx):\n",
    "        assert 0 <= idx <= n_interp\n",
    "        return st + (end - st) * idx / n_interp \n",
    "    \n",
    "    def process_list(start_list, end_list, n_interp, idx):\n",
    "        '''n_interp: how many in total between \n",
    "        '''\n",
    "        r_list = []\n",
    "        for i in range(len(start_list)):\n",
    "            st = start_list[i]\n",
    "            end = end_list[i]\n",
    "            if isinstance(st, tuple):\n",
    "                result = []\n",
    "                for iii in range(len(st)):\n",
    "                    result.append( process_elem( st[iii], end[iii], n_interp, idx) )\n",
    "                result = tuple(result)\n",
    "                # print(f\"Tuple with length {len(element)}: {element}\")\n",
    "            elif isinstance(st, float):\n",
    "                # Process if the element is a float\n",
    "                # print(f\"Float rounded to 2 decimal places: {round(element, 2)}\")\n",
    "                result = process_elem(st, end, n_interp, idx)\n",
    "            else:\n",
    "                raise NotImplementedError()\n",
    "            r_list.append( result )\n",
    "        return r_list\n",
    "\n",
    "    end_yaw = -260 # 100 # 40 + i_y\n",
    "    # gifs.extend([gifs[-1]]*10)\n",
    "    rot_gap = 3; # n_rot_frames = round(360 / rot_gap)\n",
    "    rot_frames = []\n",
    "\n",
    "    ## --------\n",
    "    ## --------\n",
    "    rot_pitch = - 40\n",
    "    start_list = list( k7d_camera_info[-4:] )\n",
    "    end_list = [end_yaw, rot_pitch, dist, cam_pos]\n",
    "    yaw_diff = end_list[0] - start_list[0]\n",
    "    n_interp = 36 if abs(yaw_diff) < 120 else int(abs(yaw_diff // 3)) # 12\n",
    "    print('XX:', end_list[0] - start_list[0])\n",
    "    for i_itp in range(n_interp):\n",
    "        i_y, i_p, i_d, i_cp = process_list(start_list, end_list, n_interp, i_itp)\n",
    "\n",
    "        \n",
    "        view_mat = p.computeViewMatrixFromYawPitchRoll(\n",
    "                cameraTargetPosition=i_cp, distance=i_d, # [0, 0, 0]\n",
    "                yaw=i_y, pitch=i_p, roll=0, upAxisIndex=2)\n",
    "        \n",
    "        rot_frames.append(p.getCameraImage(width=h, height=w, lightDirection=[1, 1, 1], shadow=0,\n",
    "                                                 renderer=p.ER_TINY_RENDERER,\n",
    "                                                 viewMatrix=view_mat,\n",
    "                                                #  projectionMatrix=env.proj_mat,\n",
    "                                                projectionMatrix=k7d_camera_info[3],\n",
    "                                                 flags=p.ER_NO_SEGMENTATION_MASK,\n",
    "                                                 )[2])\n",
    "    ## --------\n",
    "\n",
    "    \n",
    "    \n",
    "    for i_y in range(0, 360, rot_gap): # 2 3\n",
    "        view_mat = p.computeViewMatrixFromYawPitchRoll(\n",
    "                cameraTargetPosition=cam_pos, distance=dist, # [0, 0, 0]\n",
    "                yaw=(end_yaw+i_y), pitch=rot_pitch, roll=0, upAxisIndex=2)\n",
    "        \n",
    "        rot_frames.append(p.getCameraImage(width=h, height=w, lightDirection=[1, 1, 1], shadow=0,\n",
    "                                                 renderer=p.ER_TINY_RENDERER,\n",
    "                                                 viewMatrix=view_mat,\n",
    "                                                #  projectionMatrix=env.proj_mat,\n",
    "                                                projectionMatrix=k7d_camera_info[3],\n",
    "                                                 flags=p.ER_NO_SEGMENTATION_MASK,\n",
    "                                                 )[2])\n",
    "        \n",
    "        if len(rot_frames) + len(gifs) == pfr:\n",
    "            print(f'camera {pfr}', view_mat, )\n",
    "            print(f'yaw: {(end_yaw+i_y)}')\n",
    "\n",
    "\n",
    "    # gifs = set_collision_marker(gifs, has_collision, se_valid)\n",
    "    # set_collision_frame(gifs, collision_list)\n",
    "    print(f'view_mat: {view_mat}')\n",
    "    print(f'[vis traj] num_collision_frame: {num_collision_frame}; start end valid: {se_valid}')\n",
    "    \n",
    "\n",
    "    ds.extend(  [150,] * (len(rot_frames) -1) )\n",
    "    ds.append(6000)\n",
    "    gifs.extend(rot_frames)\n",
    "\n",
    "\n",
    "    # ds[-n_rot_frames:] = [200,] * n_rot_frames # slow down, a list\n",
    "    ds[0] = 1200\n",
    "    assert len(ds) == len(gifs)\n",
    "    vis_dict = dict(ncoll_frame=num_collision_frame, se_valid=se_valid)\n",
    "    env.unload_env()\n",
    "\n",
    "    return gifs, ds, vis_dict\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "import base64, datetime, pickle\n",
    "from diffuser.utils import KukaRenderer\n",
    "import diffuser.utils.luo_utils as xxx; reload(xxx)\n",
    "from diffuser.utils.luo_utils import save_gif_ethucy\n",
    "from pb_diff_envs.utils.kuka_utils_luo import visualize_kuka_traj_luo\n",
    "\n",
    "\n",
    "rd = KukaRenderer(train_env_list, is_eval=True)\n",
    "rootdir = './visualization/'\n",
    "os.makedirs(rootdir, exist_ok=True)\n",
    "gifname = f'{rootdir}/kuka7d_base.gif'\n",
    "\n",
    "\n",
    "# traj_idx = suc_idxs[0] # which path\n",
    "# hor_idx = traj_idx // bs\n",
    "# path_idx = traj_idx % bs # which path\n",
    "path_idx = 10\n",
    "\n",
    "## unnm_suc_traj list[np.ndarray]: (n_trajs, H, dim)\n",
    "gifs, ds, vis_dict = visualize_kuka_traj_luo_web( env_single, unnm_suc_traj[path_idx] )\n",
    "\n",
    "\n",
    "save_gif_ethucy( gifs, gifname, duration=ds) # duration=100)\n",
    "b64 = base64.b64encode(open(gifname, 'rb').read()).decode('ascii')\n",
    "display(HTML(f'<img src=\"data:image/gif;base64,{b64}\" />'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
